from transformers import AutoTokenizer, AutoModelForMaskedLM

# 加载 BERT 模型和分词器
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
model = AutoModelForMaskedLM.from_pretrained("google-bert/bert-base-uncased")

def predict_masked_word(text):
    # 使用 tokenizer.encode_plus 处理输入文本，并转换为 PyTorch 张量
    inputs = tokenizer.encode_plus(text, return_tensors='pt')

    # 获取输入的 input_ids
    input_ids = inputs['input_ids']

    try:
        # 确认 [MASK] 标记的位置
        masked_index = input_ids[0].tolist().index(tokenizer.mask_token_id)
    except ValueError:
        raise ValueError("输入文本中没有找到 [MASK] 标记")

    # 进行预测
    outputs = model(**inputs)
    logits = outputs.logits

    # 获取预测的标记索引
    predicted_index = logits[0, masked_index].argmax().item()

    # 解码预测的标记
    predicted_token = tokenizer.decode([predicted_index])

    return predicted_token


# 示例输入文本
text = "Hello, this is a [MASK]."

predicted_token = predict_masked_word(text)
print(f"原始文本: {text}")
print(f"预测的标记: {predicted_token}")

# 测试其他输入文本
other_text = "You are a [MASK]."
other_predicted_token = predict_masked_word(other_text)
print(f"原始文本: {other_text}")
print(f"预测的标记: {other_predicted_token}")